"""
MPU-6050 陀螺仪/加速度计传感器驱动
用于树莓派通过I2C与MPU-6050通信
原作者: MrTijn/Tijndagamer
基于MIT许可证发布
"""

import smbus3 as smbus
import time
import numpy
import os


class IMUSensor:
    # 全局变量
    GRAVITIY_MS2 = 9.80665  # 重力加速度常量
    address = None  # I2C设备地址
    bus = None  # I2C总线对象

    # 加速度计量程修饰符
    ACCEL_SCALE_MODIFIER_2G = 16384.0
    ACCEL_SCALE_MODIFIER_4G = 8192.0
    ACCEL_SCALE_MODIFIER_8G = 4096.0
    ACCEL_SCALE_MODIFIER_16G = 2048.0

    # 陀螺仪量程修饰符
    GYRO_SCALE_MODIFIER_250DEG = 131.0
    GYRO_SCALE_MODIFIER_500DEG = 65.5
    GYRO_SCALE_MODIFIER_1000DEG = 32.8
    GYRO_SCALE_MODIFIER_2000DEG = 16.4

    # 预定义量程
    ACCEL_RANGE_2G = 0x00
    ACCEL_RANGE_4G = 0x08
    ACCEL_RANGE_8G = 0x10
    ACCEL_RANGE_16G = 0x18

    GYRO_RANGE_250DEG = 0x00
    GYRO_RANGE_500DEG = 0x08
    GYRO_RANGE_1000DEG = 0x10
    GYRO_RANGE_2000DEG = 0x18

    # MPU-6050寄存器地址
    PWR_MGMT_1 = 0x6B  # 电源管理寄存器1
    PWR_MGMT_2 = 0x6C  # 电源管理寄存器2

    ACCEL_XOUT0 = 0x3B  # 加速度计X轴数据寄存器
    ACCEL_YOUT0 = 0x3D  # 加速度计Y轴数据寄存器
    ACCEL_ZOUT0 = 0x3F  # 加速度计Z轴数据寄存器

    TEMP_OUT0 = 0x41  # 温度数据寄存器

    GYRO_XOUT0 = 0x43  # 陀螺仪X轴数据寄存器
    GYRO_YOUT0 = 0x45  # 陀螺仪Y轴数据寄存器
    GYRO_ZOUT0 = 0x47  # 陀螺仪Z轴数据寄存器

    ACCEL_CONFIG = 0x1C  # 加速度计配置寄存器
    GYRO_CONFIG = 0x1B  # 陀螺仪配置寄存器

    def __init__(self, address=0x79, bus=1):
        """初始化MPU-6050传感器

        Args:
            address: I2C设备地址(默认0x79)
            bus: I2C总线号(默认1)
        """
        try:
            if not os.path.exists(f"/dev/i2c-{bus}"):
                raise FileNotFoundError(f"I2C设备 /dev/i2c-{bus} 不存在")

            self.address = address
            self.bus = smbus.SMBus(bus)
            print(f"成功初始化I2C总线 {bus}")

            # 检查设备是否响应
            try:
                self.bus.read_byte(self.address)
                print(f"成功检测到I2C设备 0x{address:02x}")
            except Exception as e:
                raise RuntimeError(f"无法访问I2C设备 0x{address:02x}: {str(e)}")

            # 唤醒设备并设置初始配置
            self.bus.write_byte_data(self.address, self.PWR_MGMT_1, 0x00)
            time.sleep(0.1)  # 等待设备稳定
            self.set_gyro_range(self.GYRO_RANGE_250DEG)
            self.set_accel_range(self.ACCEL_RANGE_2G)
            print("成功初始化MPU6050")

        except Exception as e:
            print(f"IMU传感器初始化失败: {str(e)}")
            raise

    def read_i2c_word(self, register):
        """读取两个I2C寄存器并合并数据

        Args:
            register: 要读取的第一个寄存器地址

        Returns:
            合并后的数据值
        """
        try:
            # 读取寄存器数据
            high = self.bus.read_byte_data(self.address, register)
            low = self.bus.read_byte_data(self.address, register + 1)

            value = (high << 8) + low

            if value >= 0x8000:
                return -((65535 - value) + 1)
            else:
                return value
        except Exception as e:
            print(f"读取I2C数据失败: {str(e)}")
            raise

    def get_temp(self):
        """读取MPU-6050板载温度传感器的温度

        Returns:
            温度值(摄氏度)
        """
        try:
            raw_temp = self.read_i2c_word(self.TEMP_OUT0)
            actual_temp = (raw_temp / 340.0) + 36.53
            return actual_temp
        except Exception as e:
            print(f"读取温度数据失败: {str(e)}")
            raise

    def set_accel_range(self, accel_range):
        """设置加速度计量程

        Args:
            accel_range: 要设置的量程值(建议使用预定义量程)
        """
        try:
            self.bus.write_byte_data(self.address, self.ACCEL_CONFIG, 0x00)
            self.bus.write_byte_data(self.address, self.ACCEL_CONFIG, accel_range)
            print(f"成功设置加速度计量程: {accel_range}")
        except Exception as e:
            print(f"设置加速度计量程失败: {str(e)}")
            raise

    def read_accel_range(self, raw=False):
        """读取当前加速度计量程设置

        Args:
            raw: 是否返回原始寄存器值

        Returns:
            如果raw=True,返回原始寄存器值
            如果raw=False,返回整数:-1,2,4,8或16
        """
        try:
            raw_data = self.bus.read_byte_data(self.address, self.ACCEL_CONFIG)

            if raw is True:
                return raw_data
            elif raw is False:
                if raw_data == self.ACCEL_RANGE_2G:
                    return 2
                elif raw_data == self.ACCEL_RANGE_4G:
                    return 4
                elif raw_data == self.ACCEL_RANGE_8G:
                    return 8
                elif raw_data == self.ACCEL_RANGE_16G:
                    return 16
                else:
                    return -1
        except Exception as e:
            print(f"读取加速度计量程失败: {str(e)}")
            raise

    def get_accel_data(self, g=False):
        """获取加速度计X、Y、Z轴数据

        Args:
            g: 是否以重力加速度g为单位

        Returns:
            包含测量结果的字典
        """
        try:
            x = self.read_i2c_word(self.ACCEL_XOUT0)
            y = self.read_i2c_word(self.ACCEL_YOUT0)
            z = self.read_i2c_word(self.ACCEL_ZOUT0)

            accel_scale_modifier = None
            accel_range = self.read_accel_range(True)

            if accel_range == self.ACCEL_RANGE_2G:
                accel_scale_modifier = self.ACCEL_SCALE_MODIFIER_2G
            elif accel_range == self.ACCEL_RANGE_4G:
                accel_scale_modifier = self.ACCEL_SCALE_MODIFIER_4G
            elif accel_range == self.ACCEL_RANGE_8G:
                accel_scale_modifier = self.ACCEL_SCALE_MODIFIER_8G
            elif accel_range == self.ACCEL_RANGE_16G:
                accel_scale_modifier = self.ACCEL_SCALE_MODIFIER_16G
            else:
                print("未知量程 - 使用默认2G量程")
                accel_scale_modifier = self.ACCEL_SCALE_MODIFIER_2G

            x = x / accel_scale_modifier
            y = y / accel_scale_modifier
            z = z / accel_scale_modifier

            if g is True:
                return {"x": x, "y": y, "z": z}
            elif g is False:
                x = x * self.GRAVITIY_MS2
                y = y * self.GRAVITIY_MS2
                z = z * self.GRAVITIY_MS2
                return {"x": x, "y": y, "z": z}
        except Exception as e:
            print(f"读取加速度计数据失败: {str(e)}")
            raise

    def set_gyro_range(self, gyro_range):
        """设置陀螺仪量程

        Args:
            gyro_range: 要设置的量程值(建议使用预定义量程)
        """
        try:
            self.bus.write_byte_data(self.address, self.GYRO_CONFIG, 0x00)
            self.bus.write_byte_data(self.address, self.GYRO_CONFIG, gyro_range)
            print(f"成功设置陀螺仪量程: {gyro_range}")
        except Exception as e:
            print(f"设置陀螺仪量程失败: {str(e)}")
            raise

    def read_gyro_range(self, raw=False):
        """读取当前陀螺仪量程设置

        Args:
            raw: 是否返回原始寄存器值

        Returns:
            如果raw=True,返回原始寄存器值
            如果raw=False,返回250,500,1000,2000或-1
        """
        try:
            raw_data = self.bus.read_byte_data(self.address, self.GYRO_CONFIG)

            if raw is True:
                return raw_data
            elif raw is False:
                if raw_data == self.GYRO_RANGE_250DEG:
                    return 250
                elif raw_data == self.GYRO_RANGE_500DEG:
                    return 500
                elif raw_data == self.GYRO_RANGE_1000DEG:
                    return 1000
                elif raw_data == self.GYRO_RANGE_2000DEG:
                    return 2000
                else:
                    return -1
        except Exception as e:
            print(f"读取陀螺仪量程失败: {str(e)}")
            raise

    def get_gyro_data(self):
        """获取陀螺仪X、Y、Z轴数据

        Returns:
            包含测量结果的字典
        """
        try:
            x = self.read_i2c_word(self.GYRO_XOUT0)
            y = self.read_i2c_word(self.GYRO_YOUT0)
            z = self.read_i2c_word(self.GYRO_ZOUT0)

            gyro_scale_modifier = None
            gyro_range = self.read_gyro_range(True)

            if gyro_range == self.GYRO_RANGE_250DEG:
                gyro_scale_modifier = self.GYRO_SCALE_MODIFIER_250DEG
            elif gyro_range == self.GYRO_RANGE_500DEG:
                gyro_scale_modifier = self.GYRO_SCALE_MODIFIER_500DEG
            elif gyro_range == self.GYRO_RANGE_1000DEG:
                gyro_scale_modifier = self.GYRO_SCALE_MODIFIER_1000DEG
            elif gyro_range == self.GYRO_RANGE_2000DEG:
                gyro_scale_modifier = self.GYRO_SCALE_MODIFIER_2000DEG
            else:
                print("未知量程 - 使用默认250度/秒量程")
                gyro_scale_modifier = self.GYRO_SCALE_MODIFIER_250DEG

            x = x / gyro_scale_modifier
            y = y / gyro_scale_modifier
            z = z / gyro_scale_modifier

            return {"x": x, "y": y, "z": z}
        except Exception as e:
            print(f"读取陀螺仪数据失败: {str(e)}")
            raise

    def get_all_data(self):
        """读取并返回所有可用数据

        Returns:
            包含加速度计、陀螺仪和温度数据的列表
        """
        try:
            temp = self.get_temp()
            accel = self.get_accel_data()
            gyro = self.get_gyro_data()
            return [accel, gyro, temp]
        except Exception as e:
            print(f"读取所有数据失败: {str(e)}")
            raise


if __name__ == "__main__":
    # 测试代码
    mpu = IMUSensor(0x79)
    mpu.set_gyro_range(mpu.GYRO_RANGE_2000DEG)

    while True:
        # 获取加速度计数据
        accel_data = mpu.get_accel_data()
        print("加速度计数据:")
        print(f"X: {accel_data['x']:.2f} m/s²")
        print(f"Y: {accel_data['y']:.2f} m/s²")
        print(f"Z: {accel_data['z']:.2f} m/s²")
        print()

        # 获取陀螺仪数据
        gyro_data = mpu.get_gyro_data()
        print("陀螺仪数据:")
        print(f"X: {gyro_data['x']:.2f} °/s")
        print(f"Y: {gyro_data['y']:.2f} °/s")
        print(f"Z: {gyro_data['z']:.2f} °/s")
        print()

        # 获取温度数据
        temp = mpu.get_temp()
        print(f"温度: {temp:.2f}°C")
        print("-" * 30)

        time.sleep(1)
